#include <iostream>
#include <vector>
#include <memory>
#include <thread>
#include <mutex>
#include <queue>
#include <functional>
#include <condition_variable>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <cstring>
#include <atomic>

class ThreadPool {
public:
    explicit ThreadPool(size_t threads) : stop(false) {
        for(size_t i = 0; i < threads; ++i) {
            workers.emplace_back([this] {
                while(true) {
                    std::function<void()> task;
                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex);
                        this->condition.wait(lock, [this] {
                            return this->stop || !this->tasks.empty();
                        });
                        if(this->stop && this->tasks.empty())
                            return;
                        task = std::move(this->tasks.front());
                        this->tasks.pop();
                    }
                    task();
                }
            });
        }
    }

    template<class F>
    void enqueue(F&& f) {
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            tasks.emplace(std::forward<F>(f));
        }
        condition.notify_one();
    }

    ~ThreadPool() {
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            stop = true;
        }
        condition.notify_all();
        for(std::thread &worker: workers)
            worker.join();
    }

private:
    std::vector<std::thread> workers;
    std::queue<std::function<void()>> tasks;
    std::mutex queue_mutex;
    std::condition_variable condition;
    bool stop;
};

class EpollServer {
public:
    EpollServer(int port, size_t thread_pool_size) 
        : port(port), 
          thread_pool(thread_pool_size),
          running(false) {}

    ~EpollServer() {
        stop();
    }

    void start() {
        running = true;
        
        // 创建epoll实例
        epoll_fd = epoll_create1(0);
        if(epoll_fd == -1) {
            throw std::runtime_error("Failed to create epoll instance");
        }

        // 创建监听socket
        server_fd = socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0);
        if(server_fd == -1) {
            throw std::runtime_error("Failed to create socket");
        }

        // 设置SO_REUSEADDR
        int opt = 1;
        setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

        // 绑定地址
        sockaddr_in server_addr{};
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(port);

        if(bind(server_fd, (sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
            throw std::runtime_error("Failed to bind socket");
        }

        // 开始监听
        if(listen(server_fd, SOMAXCONN) == -1) {
            throw std::runtime_error("Failed to listen on socket");
        }

        // 添加server_fd到epoll
        epoll_event event{};
        event.events = EPOLLIN | EPOLLET;  // 边缘触发模式
        event.data.fd = server_fd;
        if(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, server_fd, &event) == -1) {
            throw std::runtime_error("Failed to add server fd to epoll");
        }

        std::cout << "Server started on port " << port << std::endl;

        // 事件循环
        event_loop();
    }

    void stop() {
        if(running) {
            running = false;
            close(server_fd);
            close(epoll_fd);
        }
    }

private:
    void event_loop() {
        const int MAX_EVENTS = 640;
        epoll_event events[MAX_EVENTS];

        while(running) {
            int num_events = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
            if(num_events == -1) {
                if(errno == EINTR) continue;
                throw std::runtime_error("epoll_wait failed");
            }

            for(int i = 0; i < num_events; ++i) {
                if(events[i].data.fd == server_fd) {
                    handle_accept();
                } else {
                    if(events[i].events & EPOLLIN) {
                        thread_pool.enqueue([this, fd = events[i].data.fd] {
                            handle_client(fd);
                        });
                    }
                }
            }
        }
    }

    void handle_accept() {
        sockaddr_in client_addr{};
        socklen_t client_len = sizeof(client_addr);
        int client_fd;

        // 边缘触发需要接受所有连接
        while((client_fd = accept4(server_fd, (sockaddr*)&client_addr, 
                                 &client_len, SOCK_NONBLOCK)) > 0) {
            std::cout << "New connection from " 
                      << inet_ntoa(client_addr.sin_addr) << ":"
                      << ntohs(client_addr.sin_port) << std::endl;

            // 添加客户端到epoll
            epoll_event event{};
            event.events = EPOLLIN | EPOLLET | EPOLLRDHUP;
            event.data.fd = client_fd;
            if(epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &event) == -1) {
                std::cerr << "Failed to add client fd to epoll" << std::endl;
                close(client_fd);
            }
        }

        if(client_fd == -1 && errno != EAGAIN && errno != EWOULDBLOCK) {
            std::cerr << "accept error" << std::endl;
        }
    }

    void handle_client(int fd) {
        char buffer[4096];
        ssize_t bytes_read;

        // 边缘触发必须读取所有数据
        while((bytes_read = read(fd, buffer, sizeof(buffer))) > 0) {
            // 简单回显服务
            if(write(fd, buffer, bytes_read) != bytes_read) {
                std::cerr << "Write error" << std::endl;
                break;
            }
        }

        if(bytes_read == 0 || (bytes_read == -1 && errno != EAGAIN)) {
            // 连接关闭
            std::cout << "Connection closed on fd " << fd << std::endl;
            epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, nullptr);
            close(fd);
        }
    }

    int port;
    int server_fd;
    int epoll_fd;
    std::atomic<bool> running;
    ThreadPool thread_pool;
};

int main(int argc, char* argv[]) {

    std::cout<<"hello server"<<std::endl;
    try {
        EpollServer server(std::stoi("8080"), std::thread::hardware_concurrency());
        server.start();
    } catch(const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}